﻿#include "IdCardOcr.h"

IdCardOcr::IdCardOcr(string det_model_dir, string cls_model_dir, string rec_model_dir, string rec_label_file, string device = "CPU") {
    this->det_model_dir = det_model_dir;
    this->cls_model_dir = cls_model_dir;
    this->rec_model_dir = rec_model_dir;
    this->rec_label_file = rec_label_file;
    if (device == "GPU" || device == "gpu")
        option.UseGpu();
    if (device == "CPU" || device == "cpu")
        option.UseCpu();
    init();
}

IdCardOcr::IdCardOcr(string det_model_dir, string rec_model_dir, string rec_label_file, string device = "CPU") {
    this->det_model_dir = det_model_dir;
    this->rec_model_dir = rec_model_dir;
    this->rec_label_file = rec_label_file;
    if (device == "GPU" || device == "gpu")
        option.UseGpu();
    if (device == "CPU" || device == "cpu")
        option.UseCpu();
    init();
}

void IdCardOcr::init()
{
    auto det_model_file = det_model_dir + sep + "inference.pdmodel";
    auto det_params_file = det_model_dir + sep + "inference.pdiparams";
    det_model = fastdeploy::vision::ocr::DBDetector(det_model_file, det_params_file, option);
    assert(det_model.Initialized());
    auto rec_model_file = rec_model_dir + sep + "inference.pdmodel";
    auto rec_params_file = rec_model_dir + sep + "inference.pdiparams";
    rec_model = fastdeploy::vision::ocr::Recognizer(rec_model_file, rec_params_file, rec_label_file, option);
    assert(rec_model.Initialized());
    if (cls_model_dir.size() != 0) {
        auto cls_model_file = cls_model_dir + sep + "inference.pdmodel";
        auto cls_params_file = cls_model_dir + sep + "inference.pdiparams";
        cls_model = fastdeploy::vision::ocr::Classifier(cls_model_file, cls_params_file, option);
        assert(cls_model.Initialized());
        model = new fastdeploy::pipeline::PPOCRv3(&det_model, &cls_model, &rec_model);
    }
    else {
        model = new fastdeploy::pipeline::PPOCRv3(&det_model, &rec_model);
    }
}

string IdCardOcr::infer(Mat& im) {
    
    auto im_bak = im.clone();
    cvtColor(im_bak, im_bak, COLOR_BGR2GRAY);
    cvtColor(im_bak, im_bak, COLOR_GRAY2BGR);

    fastdeploy::vision::OCRResult result;
    if (!model->Predict(&im_bak, &result)) {
        cerr << "Failed to predict." << endl;
        return "Failed to predict.";
    }

    int resLen = result.text.size();
    for (int i = 0; i < resLen; i++) {
        if (result.rec_scores[i] >= 0.6 && result.text[i].size() >= 18) {
            int IDFlag = 0;
            string resText;
            for (char c : result.text[i]) {
                if (c == 88 || c == 120 || (c >= 48 && c <= 57)) {
                    IDFlag++;
                    resText.push_back(c);
                }
                else IDFlag = 0;
            }
            qDebug() << QString::fromStdString(resText);

            if (IDFlag == 18) return resText;
        }
    }
    return "";
}
